#!/usr/bin/python

####################################
## A base function for plots (2D) ##
####################################

#######################################################      
## Set PyRoot Environment and PlotConfig Path first ###
#######################################################
import sys
ROOTSYS = '/afs/atlas.umich.edu/opt/root/lib'
sys.path.append(ROOTSYS)
PlotConfig='../'
PlotConfig2='.'
sys.path.append(PlotConfig)
sys.path.append(PlotConfig2)

#####################
## Import Module  ###
#####################
import array
import os
import glob
from math import sqrt,fabs,sin
from ROOT import TFile,TTree,TChain,TBranch,TH1,TH1F,TH2,TH2F,TList
from ROOT import TLorentzVector,TGraphAsymmErrors,TMath
from ROOT import THStack,TCanvas,TLegend,TColor,TPaveText,TPad
from ROOT import gStyle,gDirectory
from ROOT import Double

def plot_template_2D(histodir='NOM',filetag='',hname='',oname='',configfile='plot_configure',hrange='',hxt='',hyt='',hrebinx=1,hrebiny=1,xmax=-999,xmin=-999,ymax=-999,ymin=-999,sxt=0.06,syt=0.06,offxt=1.2,offyt=1.2,showATLAS='',verbose=-1,**kw):

    '''
    # By default, 2F histogram would be drawn in 'BOX' mode
    flietag => only include histogram files containing certain tags; 
                in most cases, don't need to set this
                Note that all the histogram should be in ../Result/XXX/
    hname   => the histogram name; to find all the names, simple use 
                "python getlist_histo.py" (loo at histo.list then)
    oname   => the output picture name; if not set, will use the histogram name
    configfile => Which config file to use, default value is "plot_configure"
    hrange  => A string indicating the ORIGINAL histogram ranges: "nbinsx,xmin,xmax,nbinsy,ymin,ymax"; 
                e.g.  "1000,-10,10",1000,-10,10"
               If you want the program to read the histogram range automatically from input histo, 
               do not set this parameter
    hxt     => The title of x axis
    hyt     => The title of y axis    
    hrebinx  => How many times for histogram rebinning (x axis)
    hrebiny  => How many times for histogram rebinning (y axis)
    xmin    => the lower limit for x axis
    xmax    => the upper limit for x axis
    ymax    => The maximum of y axis; if set to -1, will use 20*ActualMaximun for logrithm plot 
                and 2*ActualMaximum for linear plot
    ymin    => The minimum of y axis    
    sxt,syt,offxt,offyt => x axis title size / y axis title size / x axis title offset / y axis title offset
    showATLAS  => If you want to show "ATLAS"/"ATLAS Preliminary" here, put the actual string here; otherwise leave it blank
    verbose => Set it to be larger than 1, if you want it output verbose message (1->output more details, 2->only dataset info)
    '''

    print 'INFO => Make Plot for histogram %s' % hname
    print ''
    
    ############################
    ## Load Atlas Plot Style ###            
    ############################
    
    from plot_style import atlas_style
    atlas_style()

    ###################
    ## Define Canvas ##
    ###################
    print 'INFO => Defining Canvas'
    
    canvas_w,canvas_h=800,600
    gStyle.SetPadLeftMargin(0.15)
    gStyle.SetPadRightMargin(0.10)
    gStyle.SetPadTopMargin(0.05)
    gStyle.SetPadBottomMargin(0.15)
    MyC=TCanvas('MyC','MyC',canvas_w,canvas_h)
    MyC.SetTicks(1,1)    
    print 'INFO => Canvas is defined as (width=%d,height=%d)' % (canvas_w,canvas_h)
    
    ##########################################
    ## Load files,Inputs and Get histograms ##
    ##########################################
    print 'INFO => Loading histograms'
    
    # Parse config file
    from plot_config import plot_config
    dict=plot_config(configfile)
    # Read variables
    IntLumi=dict['IntLumi'] # Integrated Luminosity,pb-1
    MCScale=dict['MCScale'] # Global MCScale Factor
    Dataset=dict['Datasets'] # The Dataset contribution, data/sig/zz/wz...

    # Get the histogram range info
    print 'xmin=%3d, xmax=%3d, ymin=%3d, ymax=%3d' % (xmin, xmax, ymin, ymax)
    histo_range=hrange.split(',')
    histo_all={}
    histo_all_temp={}
    histo_key={}
    histo_xnbin, histo_xmin, histo_xmax, histo_xwidth = 0, 0., 0., 0.
    histo_ynbin, histo_ymin, histo_ymax, histo_ywidth = 0, 0., 0., 0.
    temp_xnbin, temp_xmin, temp_xmax, temp_xwidth, temp_xbinoffset = 0, 0., 0., 0., 0
    temp_ynbin, temp_ymin, temp_ymax, temp_ywidth, temp_ybinoffset = 0, 0., 0., 0., 0
    if len(histo_range)==6: # if hrange is specified
        histo_xnbin, histo_xmin, histo_xmax, histo_xwidth= int(histo_range[0]),float(histo_range[1]),float(histo_range[2]), \
                                                        (float(histo_range[2])-float(histo_range[1]))/int(histo_range[0])
        histo_ynbin, histo_ymin, histo_ymax, histo_ywidth= int(histo_range[3]),float(histo_range[4]),float(histo_range[5]), \
                                                        (float(histo_range[5])-float(histo_range[4]))/int(histo_range[3])
    else: # if hrange is not specified, grab a histogram file to extract the histogram range info
        for dataset in Dataset:
            for key in dict[dataset].keys():
                if key=='Color' or key=='Corr' or key=='OverLay': continue
                # get list of folders for this dataset or process
                list_folders=[]
                if 'data' in key: list_folders.append(key)
                for folder in list_folders:
                    # Read histo file
                    if filetag: filename=glob.glob('%s/*%s*root' % (folder,filetag))
                    else: filename=glob.glob('%s/*root' % (folder))
                    file=TFile(filename[0])
                    if histodir!=' ': hdir=gDirectory.Get(histodir)
                    else: hdir=gDirectory
                    histo_temp=hdir.Get(hname)
                    if not histo_temp:
                        print 'ERROR ==> In %s, no histogram %s exist!' % (filename,hname)
                        continue
                    histo_xnbin=histo_temp.GetNbinsX()
                    histo_xmin=histo_temp.GetXaxis().GetXmin()
                    histo_xmax=histo_temp.GetXaxis().GetXmax()
                    histo_xwidth=histo_temp.GetXaxis().GetBinWidth(1)
                    histo_ynbin=histo_temp.GetNbinsY()
                    histo_ymin=histo_temp.GetYaxis().GetXmin()
                    histo_ymax=histo_temp.GetYaxis().GetXmax()
                    histo_ywidth=histo_temp.GetYaxis().GetBinWidth(1)
                    file.Close()
                    print 'xnbin=%3d, xmin=%3d, xmax=%3d, xwidth=%3d, ynbin=%3d, ymin=%3d, ymax=%3d, ywidth=%3d' % (histo_xnbin,histo_xmin, histo_xmax, histo_xwidth, histo_ynbin, histo_ymin, histo_ymax, histo_ywidth)
                    break
                break
            break
    # If xmin and xmax are specified, need to modify the histogram range        
    if xmin==-999 and xmax==-999:
        temp_xnbin, temp_xmin, temp_xmax, temp_xwidth, temp_xbinoffset = histo_xnbin, histo_xmin, histo_xmax, histo_xwidth, 0
        temp_ynbin, temp_ymin, temp_ymax, temp_ywidth, temp_ybinoffset = histo_ynbin, histo_ymin, histo_ymax, histo_ywidth, 0
    else:
        if xmin!=-999: temp_xmin=xmin
        else: temp_xmin=histo_xmin
        if xmax!=-999: temp_xmax=xmax
        else: temp_xmax=histo_xmax
        temp_xnbin=int((temp_xmax-temp_xmin)/histo_xwidth)
        temp_xbinoffset=int((temp_xmin-histo_xmin)/histo_xwidth)
        temp_xwidth = histo_xwidth
        if ymin!=-999: temp_ymin=ymin
        else: temp_ymin=histo_ymin
        if ymax!=-999: temp_ymax=ymax
        else: temp_ymax=histo_ymax
        temp_ynbin=int((temp_ymax-temp_ymin)/histo_ywidth)
        temp_ybinoffset=int((temp_ymin-histo_ymin)/histo_ywidth)
        temp_ywidth = histo_ywidth
        print 'temp_xnbin=%3d, temp_xmin=%3d, temp_xmax=%3d, temp_ynbin=%3d, temp_ymin=%3d, temp_ymax=%3d'%(temp_xnbin,temp_xmin,temp_xmax,temp_ynbin,temp_ymin,temp_ymax)
    # Define histograms for further usage  
    for dataset in Dataset:
        print dataset
        histo_all[dataset]=TH2F(dataset,dataset,temp_xnbin,temp_xmin,temp_xmax,temp_ynbin,temp_ymin,temp_ymax)
        histo_all[dataset].Sumw2()    
       
    for dataset in Dataset:
        histo_key[dataset]={}
        for key in dict[dataset].keys():
            if key=='Color' or key=='Corr' or key=='OverLay': continue
            histo_key[dataset][key]=TH2F('%s_%s' % (dataset,key),'%s_%s' % (dataset,key),histo_xnbin,histo_xmin,histo_xmax, \
                                                                                        histo_ynbin,histo_ymin,histo_ymax)
            histo_key[dataset][key].Sumw2()
    # loop all files and merge histograms (normalization will be done for MC)
    for dataset in Dataset:
        print 'INFO ==> Looking at %s' % dataset
        this_corr=dict[dataset]['Corr']
        for key in dict[dataset].keys():
            if key=='Color' or key=='Corr' or key=='OverLay': continue
            # get list of folders for this dataset or process
            if verbose>=2: print 'INFO ==> KEY = %s' % key
            list_folders=[]
            if 'data' in key: list_folders.append(key)
            else: 
                for folder in dict[dataset][key].keys(): list_folders.append(folder)
            # loop all the folders one by one and merge histograms
            xs, k, fe, entry = 0., 0., 0., 0.
            for folder in list_folders:
                # Read histo file
                if filetag: filename=glob.glob('%s/*%s*root' % (folder,filetag))
                else: filename=glob.glob('%s/*root' % (folder))
                file=TFile(filename[0])
                # Get histogram
                if histodir!=' ': hdir=gDirectory.Get(histodir)
                else: hdir=gDirectory
                histo_temp=hdir.Get(hname)
                # Check if the histogram valid or not
                if not histo_temp:
                    print 'ERROR ==> In %s, no histogram %s exist!' % (filename,hname)
                    continue
                # For MC, mormalize histogram and then combine
                if 'mc' in folder:
                    xs=float(dict[dataset][key][folder]['Xsection(pb)'])
                    k=float(dict[dataset][key][folder]['k-factor'])
                    fe=float(dict[dataset][key][folder]['FilterEff(gen)'])
                    entry+=float(dict[dataset][key][folder]['Nevts(Weight)'])
                    if verbose!=-1: print 'INFO ==> Add histograms from %s' % folder
                    histo_key[dataset][key].Add(histo_temp)
                # For Data, just combine histogram inside one dataset
                if 'data' in key:
                    if verbose>=1: print 'INFO ==> Add histograms from %s' % folder
                    histo_all[dataset].Add(histo_temp)
                    # Close histo file
                file.Close()
            # If MC contains several different subsamples in one MCID, one needs to sum 
            # the entries to normalize MC histograms, scale factor = 
            # xs*kfactor*fe*Lumi*MCScale*dict[dataset][Corr]/entry
            if not 'data' in key:
                mc_scale=xs*k*fe*IntLumi*1000.0*MCScale*this_corr/entry
                histo_key[dataset][key].Scale(mc_scale)
                histo_all[dataset].Add(histo_key[dataset][key])
                if verbose!=-1: print 'INFO ==> Normalization done for %s -> %.6f (tot entry: %.6f)' % (key, mc_scale, entry)

        # If xmin and xmax specified, try to rebuild the histogram using according range
                
    ####################
    ## Begin to plot ###
    ####################
    from plot_style import color_convert

    # Save histograms to root file
    histo_all_temp={}
    for dataset in Dataset:
        histo_all_temp[dataset]=histo_all[dataset].Clone()

    # Set Histogram Fill style
    gStyle.SetHistFillStyle(1001)

    # Set Attributes for histograms:
    # 1) Rebin
    # 2) Set Colors
    
    for dataset in Dataset:
        # Rebin
        histo_all[dataset].Rebin2D(hrebinx,hrebiny)
        # Set Color
        color_str=dict[dataset]['Color']
        color=color_convert(color_str)
        if color==-1:
                print 'ERROR ==> Can not parse the color: %s' % dict[dataset]['Color']
                sys.exit(1)
        print 'INFO ==> Set Color %s for dataset:%s' % (color_str,dataset)
        histo_all[dataset].SetFillColor(color)

            
    ###############
    ## Plotting ###
    ###############
    MyC.SetGrid(1,1)
    print  'INFO => Now make plots'
    
    # Set Plot attibute for stacked histogram and Draw the MC part
            
    # Set drawing attributes for histograms
    # Set title and its size and offset      
    
    # Draw Data
               
    ###############
    ## Add Text ###
    ###############
        

    ##################
    ## Print Plots ###
    ##################

    # print the canvas
    # output to root file
    resultRoot=TFile(hname+'.root', 'RECREATE')
    for dataset in Dataset:
        dataset_name=dataset
        if '->' in dataset:
            dataset_name=dataset_name.replace('->', '_')
        if '/' in dataset:
            dataset_name=dataset_name.replace('/', '_')
        if '+' in dataset:
            dataset_name=dataset_name.replace('+', '_')
        if '-' in dataset:
            dataset_name=dataset_name.replace('-', '_')
        histo_all_temp[dataset].SetName(dataset_name)
        histo_all_temp[dataset].SetTitle(dataset_name)
        histo_all_temp[dataset].Write()

    #####################################################
    ## Clear the histograms and canvas after plotting ###
    #####################################################
  
    for dataset in histo_all.keys():
        histo_all[dataset].SetDirectory(0)
        TH2.AddDirectory(0)
    MyC.Close()
